#include "kernel_operator.h"
using namespace AscendC;

#define BUFFER_NUM 1
#define BLOCK_SIZE  32
#define COMPUTE_ALIGN  256
template <typename T>
class kernelHeaviside
{
public:
    __aicore__ inline kernelHeaviside(){}
    __aicore__ inline void Init(GM_ADDR input, GM_ADDR values, GM_ADDR out, uint32_t iter, uint32_t coreDataNum, uint32_t tail, TPipe *pipeIn){
        this->pipe = pipeIn;
        int coreIdx = GetBlockIdx();
        int globalIndex;
        int bigCoreDataNum = coreDataNum + 1;
        uint8_t nBigCore =tail;
        globalIndex = bigCoreDataNum * coreIdx;
        if (coreIdx < nBigCore)
        {
            this->blockNum = bigCoreDataNum;
        }
        else{
            this->blockNum = coreDataNum;
            globalIndex -= (coreIdx - nBigCore);
        }
        this->blockElemCount = BLOCK_SIZE / sizeof(T);
        int32_t totalCount = blockElemCount * this->blockNum;
        // printf("core:%d globalIndex: %d totalCount %d\n",coreIdx,globalIndex,totalCount);

        indexGm.SetGlobalBuffer((__gm__ T *)input + globalIndex * blockElemCount, totalCount);
        valuesGm.SetGlobalBuffer((__gm__ T *)values + globalIndex * blockElemCount, totalCount);
        outputGm.SetGlobalBuffer((__gm__ T *)out + globalIndex * blockElemCount, totalCount);
        if (totalCount < iter * blockElemCount)
        {
            this->loop = 1;
            this->tailCount = blockNum;
            iterCount = blockNum;
            bufferElem = ((totalCount * sizeof(T) + COMPUTE_ALIGN-1) /  COMPUTE_ALIGN) * COMPUTE_ALIGN / sizeof(T);                      
        }
        else
        {
            iterCount = iter;
            this->loop = (blockNum + iterCount - 1) / iterCount;
            this->tailCount = blockNum % iterCount;
            bufferElem = ((iter * blockElemCount * sizeof(T)+COMPUTE_ALIGN-1) / COMPUTE_ALIGN) * COMPUTE_ALIGN / sizeof(T); 
        }
        pipe->InitBuffer(inputQueue, BUFFER_NUM, bufferElem * sizeof(T));
        pipe->InitBuffer(valuesQueue, BUFFER_NUM, bufferElem*sizeof(T));
        pipe->InitBuffer(outQueue, BUFFER_NUM, bufferElem*sizeof(T));
        pipe->InitBuffer(bitBuf,  bufferElem / 8 );
    }
    __aicore__ inline void process(){
        curElem = iterCount * blockElemCount;
        this->preIter = curElem;
        round = ((curElem * sizeof(T) + COMPUTE_ALIGN - 1) / COMPUTE_ALIGN) * COMPUTE_ALIGN / sizeof(T);
        for (int i = 0; i < loop - 1;i++)
        {
            copyIn(i);
            compute(i);
            copyOut(i);
        }
        curElem = tailCount * blockElemCount;
        round = ((curElem * sizeof(T) + COMPUTE_ALIGN - 1) / COMPUTE_ALIGN) * COMPUTE_ALIGN / sizeof(T);
        copyIn(loop - 1);
        compute(loop - 1);
        copyOut(loop - 1);
    }
    __aicore__ inline void copyIn(int i){
        auto inputLocal = inputQueue.AllocTensor<T>();
        auto valuesLocal = valuesQueue.AllocTensor<T>();
        DataCopy(inputLocal, indexGm[preIter*i], curElem);
        DataCopy(valuesLocal, valuesGm[preIter*i], curElem);
        inputQueue.EnQue(inputLocal);
        valuesQueue.EnQue(valuesLocal);
    }
    __aicore__ inline void compute(int i){
        auto outLocal = outQueue.AllocTensor<T>();
        auto inputLocal = inputQueue.DeQue<T>();
        auto cmpRes = bitBuf.Get<uint8_t>();

        Duplicate(outLocal, (T)0, round);
        // AscendC::PipeBarrier<PIPE_V>();
        CompareScalar(cmpRes, inputLocal, (T)0, CMPMODE::LE, round);
        // printf("cmpRes %d\n", cmpRes.GetValue(2049/8));
        // AscendC::PipeBarrier<PIPE_V>();

        Select(outLocal, cmpRes, outLocal,(T)1,  SELMODE::VSEL_TENSOR_SCALAR_MODE, round);
        // printf("outLocal %f\n", outLocal.GetValue(2049));
        // AscendC::PipeBarrier<PIPE_V>();

        CompareScalar(cmpRes, inputLocal, (T)0, CMPMODE::EQ, round);
        // AscendC::PipeBarrier<PIPE_V>();

        auto valuesLocal = valuesQueue.DeQue<T>();
        Select(outLocal, cmpRes, valuesLocal, outLocal,  SELMODE::VSEL_TENSOR_TENSOR_MODE, round);
        // AscendC::PipeBarrier<PIPE_V>();

        outQueue.EnQue(outLocal);
        inputQueue.FreeTensor(inputLocal);
        valuesQueue.FreeTensor(valuesLocal);
    }
    __aicore__ inline void copyOut(int i){
        auto outLocal = outQueue.DeQue<T>();
        DataCopy(outputGm[preIter*i], outLocal, curElem);
        outQueue.FreeTensor(outLocal);
    }

private:
    TPipe *pipe;
    GlobalTensor<T> indexGm;
    GlobalTensor<T> valuesGm;
    GlobalTensor<T> outputGm;
    TQue<QuePosition::VECIN, BUFFER_NUM> inputQueue;
    TQue<QuePosition::VECIN, BUFFER_NUM> valuesQueue;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueue;
    TBuf<QuePosition::VECCALC> bitBuf;
    int32_t tailCount, blockNum, bufferElem, loop , curElem, iterCount, round, blockElemCount,preIter;
};


extern "C" __global__ __aicore__ void heaviside(GM_ADDR input, GM_ADDR values,
                                                GM_ADDR out, GM_ADDR workspace,
                                                GM_ADDR tiling) {
  GET_TILING_DATA(tiling_data, tiling);
  // TODO: user kernel impl
  TPipe pipe;
  kernelHeaviside<DTYPE_INPUT> op;
  op.Init(input, values, out, tiling_data.iter,
          tiling_data.coreDataNum, tiling_data.tail, &pipe);
  op.process();
}